import sys
parent_dir = "../pygl"
sys.path.append(parent_dir)
import time
import OpenGL.GL as gl
import glm
import imgui
import numpy as np
import pygl
from common import AppWithTrackball, AxesHelper, GridHelper, ColorMap

class FLTReader:
    def __init__(self, file_name:str) -> None:
        if ".flt" in file_name:
            header_file = file_name.replace(".flt", ".hdr")
            with open(header_file) as file:
                header_info = file.readlines()
                header_items = dict(line.upper().strip().split() for line in header_info)
                self.ncols = int(header_items['NCOLS'])
                self.nrows = int(header_items['NROWS'])
                self.cellsize = float(header_items['CELLSIZE'])
                if header_items.get('NODATA_VALUE') is not None:
                    self.nodata_value = float(header_items['NODATA_VALUE'])
                self.byteorder = header_items['BYTEORDER']
                if header_items.get('XLLCORNER') is not None:
                    self.xllcorner = float(header_items['XLLCORNER'])
                    self.xllcenter = self.xllcorner + 0.5 * self.cellsize
                elif header_items.get('XLLCENTER')  is not None:
                    self.xllcenter = float(header_items['XLLCENTER'])
                    self.xllcorner = self.xllcenter - 0.5 * self.cellsize
                if header_items.get('YLLCORNER') is not None:
                    self.yllcorner = float(header_items['YLLCORNER'])
                    self.yllcenter = self.yllcorner + 0.5 * self.cellsize
                elif header_items.get('YLLCENTER') is not None:
                    self.yllcenter = float(header_items['YLLCENTER'])
                    self.yllcorner = self.yllcenter - 0.5 * self.cellsize
            with open(file_name, "rb") as file:
                self.dem_data = np.fromfile(file, dtype = np.float32)
                if self.byteorder != 'LSBFIRST':
                    self.dem_data.byteswap()
            self.dem_data[self.dem_data == self.nodata_value] = np.nan
            self.dem_data= self.dem_data.reshape(self.nrows, self.ncols)
            self.min = np.nanmin(self.dem_data)
            self.max = np.nanmax(self.dem_data)

class PointCloudsApp(AppWithTrackball):
    def setupVertex(self) -> None:
        gl.glEnable(gl.GL_DEPTH_TEST)
        self.background_color = glm.vec4(1.0, 1.0, 1.0, 1.0)
        self.program = pygl.ProgramVF("usecase/shaders/dem.vs", 
                                      "usecase/shaders/dem.fs")
        self.gf_file = FLTReader("dataset/DEM_Ravone.flt")
        self.num_points = self.gf_file.ncols * self.gf_file.nrows
        self.vertices = np.empty((self.num_points, 6), dtype=gl.GLfloat)
        dem_range = self.gf_file.max - self.gf_file.min
        raster_scale = 10.0
        dem_scale = 5.0
        for i in range(self.gf_file.ncols):
            for j in range(self.gf_file.nrows):
                index = i * self.gf_file.nrows + j
                self.vertices[index, 0] = i/raster_scale
                self.vertices[index, 2] = j/raster_scale
                y = self.gf_file.dem_data[j, i]
                self.vertices[index, 1] = y / dem_scale
                if not np.isnan(y):
                    color = (y - self.gf_file.min) / dem_range
                    color = ColorMap.get_color(color, ColorMap.viridis)
                    self.vertices[index, 3] = color[0]
                    self.vertices[index, 4] = color[1]
                    self.vertices[index, 5] = color[2]
        self.indices = []
        for i in range(self.gf_file.ncols - 1):
            for j in range(self.gf_file.nrows - 1):
                z1 = self.gf_file.dem_data[j, i]
                if not np.isnan(z1):
                    z3 = self.gf_file.dem_data[j + 1, i + 1]
                    if not np.isnan(z3):
                        z2 = self.gf_file.dem_data[j + 1, i]
                        if not np.isnan(z2):
                            self.indices.append(i * self.gf_file.nrows + j)
                            self.indices.append(i * self.gf_file.nrows + j + 1)
                            self.indices.append((i + 1) * self.gf_file.nrows + j + 1)
                        z2 = self.gf_file.dem_data[j, i + 1]
                        if not np.isnan(z2):
                            self.indices.append((i + 1)  * self.gf_file.nrows + j + 1)
                            self.indices.append((i + 1) * self.gf_file.nrows + j)
                            self.indices.append(i * self.gf_file.nrows + j)
        self.vertex_vbo = pygl.VertexBufferObject(self.vertices)
        self.vertex_ebo =pygl.ElementBufferObject(np.array(self.indices))
        self.vao = pygl.VertexArrayObject()
        binding_point = 0
        self.vao.setVertexBuffer(self.vertex_vbo, binding_point, 0, 
                                 6 * gl.sizeof(gl.GLfloat))
        attribute_position = pygl.VertexAttribute("position", 0, 3, 
                                                  gl.GL_FLOAT, False, 0)
        attribute_color = pygl.VertexAttribute("color",1, 3, gl.GL_FLOAT, 
                                False, 3 * gl.sizeof(gl.GLfloat))
        self.vao.setVertexAttribute(binding_point, attribute_position)
        self.vao.setVertexAttribute(binding_point, attribute_color)
        self.vao.setElementBuffer(self.vertex_ebo)
        self.model = glm.mat4(1.0)
        self.model = glm.translate(self.model, glm.vec3(-self.gf_file.ncols / 
            raster_scale / 2.0, 0, -self.gf_file.nrows / raster_scale / 2.0))
        self.camera.far = 800
        self.camera.zoom_speed = 6.0
        self.camera.reset(glm.vec3(64.0, 190.0, -20.0), glm.vec3(0.0, 0.0, 0.0), 
                    glm.vec3(0.0, 1.0, 0.0))
        self.world_axes = AxesHelper(150.0, 5.0)
        self.grid = GridHelper(150.0, 20, glm.vec3(1.0, 0.0, 0.0), 
                               glm.vec3(0.5, 0.5, 0.5))
        self.show_world_axes = True
        self.show_grid = True
        self.imgui_window_width = 220
        self.imgui_window_height = 220

    # 配置imgui界面元素
    def showImGui(self)->None:
        imgui.begin("设置")
        if imgui.button("重置轨迹球相机"):
            self.camera.reset(glm.vec3(64.0, 190.0, -20.0), glm.vec3(0.0, 0.0, 0.0), 
                    glm.vec3(0.0, 1.0, 0.0))
        _, self.show_world_axes = imgui.checkbox("显示三维世界坐标轴", 
                                                self.show_world_axes)
        _, self.show_grid = imgui.checkbox("显示XZ平面网格", 
                                                self.show_grid)
        imgui.end()   

    def render(self) -> None:
        gl.glClearNamedFramebufferfv(0, gl.GL_COLOR, 0, 
                                     self.background_color.to_tuple())
        gl.glClearNamedFramebufferfv(0, gl.GL_DEPTH, 0, [1.0]); 
        self.program.use()
        view = self.camera.getViewMatrix()
        projection = self.camera.getProjectionMatrix()
        self.program.setUniformMatrix4fv(0, self.model)
        self.program.setUniformMatrix4fv(1, view)
        self.program.setUniformMatrix4fv(2, projection)
        self.vao.bind()
        gl.glDrawElements(gl.GL_TRIANGLES, len(self.indices), gl.GL_UNSIGNED_INT, None)
        if self.show_world_axes:
            self.world_axes.updateModelMatrix
            self.world_axes.render(view, projection)
        if self.show_grid:
            self.grid.render(view, projection)

    def cleanup(self) -> None:
        self.vao.delete()
        self.vertex_vbo.delete()
        self.vertex_ebo.delete()
        self.program.delete()
        self.grid.delete()
        self.world_axes.delete()

if __name__ == "__main__":
    app = PointCloudsApp(100, 100, 1000, 600, "你好，数字高程模型App！")
    app.run()

